#include "server.h"

EchoServer::EchoServer(int port) : port_(port), listenFd_(-1), epollFd_(-1), totalConnections_(0) {}

EchoServer::~EchoServer() {
    if (listenFd_ != -1) close(listenFd_);
    if (epollFd_ != -1)  close(epollFd_);
}

// 设置socket为非阻塞模式
void EchoServer::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) {
        std::cerr << "fcntl F_GETFL failed: " << strerror(errno) << std::endl;
        return;
    }
    if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
        std::cerr << "fcntl F_SETFL O_NONBLOCK failed: " << strerror(errno) << std::endl;
    }
}

// 完整发送数据（确保所有数据都被发送）
int EchoServer::sendAll(int fd, const char* buffer, int len) {
    int total = 0;
    while (total < len) {
        int n = write(fd, buffer + total, len - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                // 发送缓冲区满，需要等待可写事件
                return -1;
            }
            return -2; // 其他错误
        }
        total += n;
    }
    return total;
}

// 初始化服务器
bool EchoServer::init() {
    // 创建监听socket
    listenFd_ = socket(AF_INET, SOCK_STREAM, 0);
    if (listenFd_ == -1) {
        std::cerr << "Create socket failed: " << strerror(errno) << std::endl;
        return false;
    }

    // 设置SO_REUSEADDR（允许端口复用）
    int opt = 1;
    if (setsockopt(listenFd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) == -1) {
        std::cerr << "Setsockopt SO_REUSEADDR failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    // 绑定IP和端口
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(port_);
    if (bind(listenFd_, (struct sockaddr*)&addr, sizeof(addr)) == -1) {
        std::cerr << "Bind failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    // 开始监听
    if (listen(listenFd_, SOMAXCONN) == -1) {
        std::cerr << "Listen failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    // 创建epoll实例
    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        close(listenFd_);
        return false;
    }

    // 将监听socket加入epoll
    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = listenFd_;
    if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, listenFd_, &ev) == -1) {
        std::cerr << "Epoll_ctl add listenFd failed: " << strerror(errno) << std::endl;
        close(epollFd_);
        close(listenFd_);
        return false;
    }

    std::cout << "Server initialized on port " << port_ << std::endl;
    return true;
}

// 启动事件循环
void EchoServer::start() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];

    while (true) {
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, -1);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            std::cerr << "Epoll wait failed: " << strerror(errno) << std::endl;
            break;
        }

        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            // 处理新连接
            if (fd == listenFd_) {
                struct sockaddr_in clientAddr;
                socklen_t clientAddrLen = sizeof(clientAddr);
                int clientFd = accept(listenFd_, (struct sockaddr*)&clientAddr, &clientAddrLen);
                if (clientFd == -1) {
                    if (errno == EAGAIN || errno == EWOULDBLOCK) continue;
                    std::cerr << "Accept failed: " << strerror(errno) << std::endl;
                    continue;
                }

                setNonBlocking(clientFd);
                totalConnections_++;
                std::cout << "New client (total: " << totalConnections_ << "): " 
                          << inet_ntoa(clientAddr.sin_addr) << ":" 
                          << ntohs(clientAddr.sin_port) << std::endl;

                struct epoll_event ev;
                ev.events = EPOLLIN;
                ev.data.fd = clientFd;
                if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, clientFd, &ev) == -1) {
                    std::cerr << "Epoll_ctl add clientFd failed: " << strerror(errno) << std::endl;
                    close(clientFd);
                }
            } 
            // 处理客户端数据
            else {
                // 处理读事件（接收数据）
                if (event & EPOLLIN) {
                    int nbytes = read(fd, buffer, BUFFER_SIZE - 1);
                    if (nbytes <= 0) {
                        if (nbytes < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) continue;
                        
                        // 连接关闭或错误
                        std::cout << "Client disconnected or error: " << fd << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        pendingData_.erase(fd);
                        continue;
                    }
                    
                    buffer[nbytes] = '\0';
                    std::cout << "Server received from fd " << fd << ": '" << buffer << "' (" << nbytes << " bytes)" << std::endl;

                    // 立即回显数据
                    int sendResult = sendAll(fd, buffer, nbytes);
                    if (sendResult == nbytes) {
                        std::cout << "Server echoed to fd " << fd << ": '" << buffer << "'" << std::endl;
                    } else if (sendResult == -1) {
                        // 发送缓冲区满，缓存数据并等待可写事件
                        pendingData_[fd] = std::string(buffer, nbytes);
                        struct epoll_event ev;
                        ev.events = EPOLLIN | EPOLLOUT;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        std::cout << "Server buffered data for fd " << fd << ", will send later" << std::endl;
                    } else {
                        std::cerr << "Server send failed for fd " << fd << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        pendingData_.erase(fd);
                    }
                }

                // 处理写事件（发送缓存的数据）
                if (event & EPOLLOUT) {
                    auto it = pendingData_.find(fd);
                    if (it == pendingData_.end()) {
                        // 没有待发送数据，恢复监听读事件
                        struct epoll_event ev;
                        ev.events = EPOLLIN;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        continue;
                    }

                    std::string& data = it->second;
                    int sendResult = sendAll(fd, data.c_str(), data.size());
                    if (sendResult == data.size()) {
                        std::cout << "Server sent buffered data to fd " << fd << ": '" << data << "'" << std::endl;
                        pendingData_.erase(it);
                        struct epoll_event ev;
                        ev.events = EPOLLIN;
                        ev.data.fd = fd;
                        epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                    } else if (sendResult == -2) {
                        std::cerr << "Server send failed for fd " << fd << std::endl;
                        pendingData_.erase(it);
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                    }
                }
            }
        }
    }
}
